import os
from flask import Flask, render_template, request, redirect, url_for
import uuid
import subprocess

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['RESULT_FOLDER'] = 'static/results'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'gif'}

# 确保文件夹存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

@app.route('/')
def index():
    # 从请求参数中获取文件名
    filename = request.args.get('filename')
    dehazed_filename = request.args.get('dehazed_filename')
    return render_template('index.html', filename=filename, dehazed_filename=dehazed_filename)

@app.route('/upload', methods=['POST'])
def upload_file():
    # 检查是否选择了文件
    if 'file' not in request.files:
        return redirect(request.url)
    file = request.files['file']
    
    # 检查是否选择了文件
    if file.filename == '':
        return redirect(url_for('index'))
    
    # 检查文件类型
    if file and allowed_file(file.filename):
        # 生成唯一文件名
        filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1]
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        # 重定向到首页，并传递文件名
        return redirect(url_for('index', filename=filename))
    
    return redirect(url_for('index'))

@app.route('/dehaze', methods=['POST'])
def dehaze_image():
    filename = request.form.get('filename')
    if not filename:
        return "No filename provided", 400
    
    input_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    
    # 确保文件存在
    if not os.path.exists(input_path):
        return "File not found", 404
    
    # 创建去雾后的文件名（保留扩展名）
    output_filename = "dehazed_" + filename
    output_path = os.path.join(app.config['RESULT_FOLDER'], output_filename)
    
    # 调用main.py处理图片
    try:
        # 检查main.py是否存在
        if not os.path.exists('main.py'):
            return "main.py not found", 500
            
        # 调用main.py
        result = subprocess.run(['python', 'main.py', input_path, output_path], 
                                capture_output=True, text=True)
        
        if result.returncode != 0:
            app.logger.error(f"main.py failed with error: {result.stderr}")
            return f"Error during dehazing: {result.stderr}", 500
    except Exception as e:
        app.logger.error(f"Exception occurred: {str(e)}")
        return f"Error during dehazing: {str(e)}", 500
    
    # 检查输出文件是否创建成功
    if not os.path.exists(output_path):
        app.logger.error(f"Output file not created: {output_path}")
        return "Output file creation failed", 500
    
    return redirect(url_for('index', filename=filename, dehazed_filename=output_filename))

if __name__ == '__main__':
    app.run()